# https://www.bilibili.com/video/BV17w411k7a9/?spm_id_from=333.1007.top_right_bar_window_custom_collection.content.click&vd_source=b664701fe1351134c3243c7903d11516
# 数据集 MINIST 里面都是一张一张的图片 
# 每张图是28*28像素(即784个格子)，灰度图，只有一个通道，每个像素点(格子)只有一个数值，
# 所以每张图(每个样本)，都有784个数字
# 每张图宏观显示都是手写的一个数字
# 所以，输出层为10个 ，要预测改2样本为数字0～9的哪一个
# 所以一个样本有784个特征点，所以全连接神经网络输入层数量为784 即x_1_1,x_1_2,...,x_784_1
import numpy as np #要用的MINIST数据集 是numpy格式
import torch
import gzip

def load_mnist_images(filename):
    """加载MNIST图像文件"""
    with gzip.open(filename, 'rb') as f:
        data = np.frombuffer(f.read(), np.uint8, offset=16)
    data = data.reshape(-1, 28, 28)
    return data

def load_mnist_labels(filename):
    """加载MNIST标签文件"""
    with gzip.open(filename, 'rb') as f:
        data = np.frombuffer(f.read(), np.uint8, offset=8)
    return data

# 加载训练集和测试集
train_images = load_mnist_images('/home/zlc/xiaomi_RL/MLP/tao_lao_shi_MLP/pytorch_MLP_2/tensorflow-tutorial-samples/mnist/data_set/train-images-idx3-ubyte.gz')
train_labels = load_mnist_labels('/home/zlc/xiaomi_RL/MLP/tao_lao_shi_MLP/pytorch_MLP_2/tensorflow-tutorial-samples/mnist/data_set/train-labels-idx1-ubyte.gz')
test_images = load_mnist_images('/home/zlc/xiaomi_RL/MLP/tao_lao_shi_MLP/pytorch_MLP_2/tensorflow-tutorial-samples/mnist/data_set/t10k-images-idx3-ubyte.gz')
test_labels = load_mnist_labels('/home/zlc/xiaomi_RL/MLP/tao_lao_shi_MLP/pytorch_MLP_2/tensorflow-tutorial-samples/mnist/data_set/t10k-labels-idx1-ubyte.gz')

# 使用字典保存，避免形状不一致的问题
mnist_data = {
    'train_x': train_images,
    'test_x': test_images,
    'train_y': train_labels,
    'test_y': test_labels
}

# 保存为npy文件
np.save('mnist.npy', mnist_data)
print("MNIST数据已保存为 mnist.npy")

#加载数据集
# 解包数据集 训练集的数据，测试集的数据，训练集的标签，测试集的标签
mnist_data = np.load('mnist.npy', allow_pickle=True).item()
train_x = mnist_data['train_x']
test_x = mnist_data['test_x']
train_y = mnist_data['train_y']
test_y = mnist_data['test_y']
print(f"训练集图像形状: {train_x.shape}")  # (60000, 28, 28) 28行28列
print(f"训练集标签形状: {train_y.shape}")   # (60000,)
print(f"测试集图像形状: {test_x.shape}")   # (10000, 28, 28)
print(f"测试集标签形状: {test_y.shape}")    # (10000,)

x_train = train_x.reshape(60000,784).astype(np.float32)   
print(f"训练集图像形状: {x_train.shape}")  # 转换成向量 784列


# np的数据类型转换成pytorch的tenser类型
x_train = torch.tensor(x_train)
y_train = torch.tensor(train_y)
print(f"*****************************")  
print(f"训练集图像: {x_train.shape}")  # 转换成向量 784列
print(f"训练集标签: {y_train.shape}")  


print(f"*******************************搭建全连接神经网络*********************************") 
import  torch.nn as nn
import torch.nn.functional as F #激活函数 交叉熵损失函数都在这个里面
# 创建神经网络模型对应的类    
class MNIST(nn.Module):# 继承nn.Module
    def __init__(self): # 成员函数：初始化方法，可以知道当前定义的类 之后实例出来的对象(我们的模型)里面可以包含有哪些属性(成员变量)和方法(成员函数)
        super().__init__() #继承父类的东西
        self.fc1 = nn.Linear(784,1000) #类的属性 第一个全连接层 1000为第二层的神经元数量  用nn里面的Linear帮我们搭建(Linear帮忙完成了wx+b的事情)
        self.fc2 = nn.Linear(1000,500)
        self.fc3 = nn.Linear(500,10)
    def forward(self, x): #前向传播函数
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x) # 不用softmax，因为后面交叉熵损失函数包含了softmax
                        # 最后一层的x通过softmax 求出p， p再去和标签y 用交叉熵求出loss
        return x

# 实例化
mnist = MNIST() #mnist这个对象 就是个神经网络模型了
y = mnist(x_train) #求出了模型的输出


print(f"*******************************求损失函数*********************************") 
loss = F.cross_entropy(y, y_train)

print(f"*******************************反向传播,求每个w的梯度值grad********************************") 
loss.backward()
lr = 0.001 #学习率


print(f"*******************************更新参数w(权重)********************************") 
with torch.no_grad(): # 上下文管理器
    for w in mnist.parameters(): # 找到每一个模型的参数w
        w -= lr * w.grad # 梯度下降 #SGD

print(f"*******************************训练 更新参数w(权重) 100次********************************") 
for i in range(100): # 训练100次 参数更新100次
    y = mnist(x_train) #求出了模型的输出
    loss = F.cross_entropy(y, y_train)
    loss.backward()
    with torch.no_grad(): # 上下文管理器
        for w in mnist.parameters(): # 找到每一个模型的参数w
            w -= lr * w.grad # 梯度下降
    mnist.zero_grad() #该函数作用：梯度清零
    # if i%5==0: #每5次迭代 就进入
    print(i,loss) # 第一个loss如果很小 说明初始化的w 初始化的很合适
